'''
    对网络的每层输出和卷积核参数进行可视化
'''
import os
import numpy as np
import cv2
from PIL import Image
import torch
from torchsummary import summary
from classification import Classification
from utils.utils import predict_input, makedir



#------------------------------------------------------#
#   visual_out_path        特征层和卷积层可视化图片保存的文件夹
#------------------------------------------------------#
visual_out_path = "net_visualize"

class NetVisualize(Classification):

    #---------------------------------------------------#
    #   获得模型信息及参数大小
    #---------------------------------------------------#
    def get_model_info(self):
        print("get_model_info:")
        device = 'cuda' if self.cuda else 'cpu'
        summary(model=self.model, input_size=(3, self.input_shape[0], self.input_shape[1]), device=device)

    #---------------------------------------------------#
    #   计算模型Parameters参数量、FLOPs浮点运算次数
    #---------------------------------------------------#
    def get_model_metrics(self):
        from thop import profile

        print("get model params and Flops")
        makedir(visual_out_path)
        x = torch.randn((1, 3, self.input_shape[0], self.input_shape[1]))
        flops, total_params = profile(self.model, inputs=(x,))
        with open(os.path.join(visual_out_path, "metrics.txt"), 'w') as f:
            f.write("Total params: %d, %.2f M\n" % (total_params, total_params / 1e6))
            f.write("Total Flops: %d Flops, %.2f GFlops" % (flops, flops / 1e9))


    #--------------------------------------------#
    #   输入Image类型的数据，提取每一层网络的输出特征层
    #   therd_size指定保存图像的大小
    #--------------------------------------------#
    def get_feature_output(self, image, therd_size):
        image = Image.open(image)
        #---------------------------------------------#
        #   对测试图片预处理
        #---------------------------------------------#
        image = predict_input(image, self.input_shape)
        #---------------------------------------------------------#
        #   添加batch_size维度
        #---------------------------------------------------------#
        image = np.expand_dims(image, 0)

        with torch.no_grad():
            x = torch.from_numpy(image)
            if self.cuda:
                x = x.cuda()
        #
        output_feature_layer = []

        # model.named_modules() 返回generator, 能够按顺序迭代网络的每一层及其子层以及名字
        # for name, module in self.model.named_modules():
        #     # 如果该层不为嵌套层,进行前向传播
        #     if (
        #             not isinstance(module, torch.nn.Sequential)
        #             and not isinstance(module, torch.nn.ModuleList)
        #             and not (module == self.model)
        #     ):
        #         print(name)
        #         print(module)

        # model._modules.items()获取网络的最外层及其名字
        # features
        # conv
        # avgpool
        # classifier
        for name, module in self.model.module._modules.items():
        # for name, module in self.model._modules.items():
            if "features" in name:
                for layer in module:
                    # print(layer)
                    with torch.no_grad():
                        x = layer(x)
                    output_feature_layer.append(x)
            elif "avgpool" in name:
                with torch.no_grad():
                    x = module(x)
                output_feature_layer.append(x)
                x = x.view(x.size(0), -1)
            else:
                with torch.no_grad():
                    x = module(x)
                if "classifier" in name:
                    x = x.view(x.size(0), x.size(1), 1, 1)
                output_feature_layer.append(x)

        # 对获得的特征层列表进行遍历
        for i, imagedata in enumerate(output_feature_layer):
            features = imagedata[0]  # 去除batch_size维度
            channels = features.shape[0]  # 得到该特征层的通道数

            dst_path = os.path.join(visual_out_path, "features", "layer_{}".format(i))
            makedir(dst_path)

            # 对特征层的每个通道遍历
            for c in range(channels):
                feature = features.data.cpu().numpy()
                feature_img = feature[c, :, :]
                feature_img = np.asarray(feature_img * 255, dtype=np.uint8)

                # 对图像进行伪彩色增强,热力图
                feature_img = cv2.applyColorMap(feature_img, cv2.COLORMAP_JET)

                # 如果图片尺寸小于therd_size, 将图像resize到therd_size尺寸
                if feature_img.shape[0] < therd_size:
                    tmp_file = os.path.join(dst_path, "channel_{}_{}.jpg".format(c, therd_size))
                    tmp_img = feature_img.copy()
                    tmp_img = cv2.resize(tmp_img, (therd_size, therd_size), interpolation=cv2.INTER_NEAREST)
                    cv2.imwrite(tmp_file, tmp_img)
                else:
                    dst_file = os.path.join(dst_path, "channel_{}.jpg".format(c))
                    cv2.imwrite(dst_file, feature_img)

    # --------------------------------------------#
    #   输入Image类型的数据，提取每一层卷积层的输出特征层
    #   therd_size指定保存图像的大小
    # --------------------------------------------#
    def get_conv_feature_output(self, image, therd_size):
        image = Image.open(image)
        # ---------------------------------------------#
        #   对测试图片预处理
        # ---------------------------------------------#
        image = predict_input(image, self.input_shape)
        # ---------------------------------------------------------#
        #   添加batch_size维度
        # ---------------------------------------------------------#
        image = np.expand_dims(image, 0)

        with torch.no_grad():
            x = torch.from_numpy(image)
            if self.cuda:
                x = x.cuda()
        #
        output_feature_layer = []

        # model.named_modules() 返回generator, 能够按顺序迭代网络的每一层及其子层以及名字
        # for name, module in self.model.named_modules():
        #     # 如果该层不为嵌套层,进行前向传播
        #     if (
        #             not isinstance(module, torch.nn.Sequential)
        #             and not isinstance(module, torch.nn.ModuleList)
        #             and not (module == self.model)
        #     ):
        #         print(name)
        #         print(module)

        # model._modules.items()获取网络的最外层及其名字
        for name, module in self.model.module._modules.items():
            # for name, module in self.model._modules.items():
            if "features" in name:
                for layer in module:

                    if isinstance(layer, torch.nn.Sequential):
                        for i in layer:
                            print(layer)
                            with torch.no_grad():
                                x = layer(x)

                    output_feature_layer.append(x)

        # 对获得的特征层列表进行遍历
        for i, imagedata in enumerate(output_feature_layer):
            features = imagedata[0]  # 去除batch_size维度
            channels = features.shape[0]  # 得到该特征层的通道数

            dst_path = os.path.join(visual_out_path, "features", "layer_{}".format(i))
            makedir(dst_path)

            # 对特征层的每个通道遍历
            for c in range(channels):
                feature = features.data.cpu().numpy()
                feature_img = feature[c, :, :]
                feature_img = np.asarray(feature_img * 255, dtype=np.uint8)

                # 对图像进行伪彩色增强,热力图
                feature_img = cv2.applyColorMap(feature_img, cv2.COLORMAP_JET)

                # 如果图片尺寸小于therd_size, 将图像resize到therd_size尺寸
                if feature_img.shape[0] < therd_size:
                    tmp_file = os.path.join(dst_path, "channel_{}_{}.jpg".format(c, therd_size))
                    tmp_img = feature_img.copy()
                    tmp_img = cv2.resize(tmp_img, (therd_size, therd_size), interpolation=cv2.INTER_NEAREST)
                    cv2.imwrite(tmp_file, tmp_img)
                else:
                    dst_file = os.path.join(dst_path, "channel_{}.jpg".format(c))
                    cv2.imwrite(dst_file, feature_img)

    #--------------------------------------------#
    #   对模型的卷积核进行可视化
    #   therd_size指定保存图像的大小
    #--------------------------------------------#
    def get_conv_visual(self, therd_size):
        # 遍历模型中每一层网络的参数
        for name, parameters in self.model.named_parameters():
            # 如果参数的维度为4,说明是卷积层
            if parameters.ndim == 4:

                dst_path = os.path.join(visual_out_path, "conv_visual", "layer_{}".format(name))
                makedir(dst_path)

                # 卷积核的个数
                kernel_number = parameters.shape[0]
                # 卷积核的通道数
                kernel_channel = parameters.shape[1]
                # 遍历卷积核的个数
                for n in range(kernel_number):
                    kernel = parameters[n, :, :, :]
                    # 遍历卷积核的通道
                    for c in range(kernel_channel):
                        kernel_image = kernel[c, :, :]

                        feature_img = np.asarray(kernel_image.cpu().detach().numpy() * 255, dtype=np.uint8)
                        # 对图像进行伪彩色增强,热力图
                        # feature_img = cv2.applyColorMap(feature_img, cv2.COLORMAP_JET)

                        # 如果图片尺寸小于therd_size, 将图像resize到therd_size尺寸
                        if feature_img.shape[0] < therd_size:
                            tmp_file = os.path.join(dst_path, "channel_{}_{}.jpg".format(c, therd_size))
                            tmp_img = feature_img.copy()
                            tmp_img = cv2.resize(tmp_img, (therd_size, therd_size), interpolation=cv2.INTER_CUBIC)
                            cv2.imwrite(tmp_file, tmp_img)
                        else:
                            dst_file = os.path.join(dst_path, "channel_{}.jpg".format(c))
                            cv2.imwrite(dst_file, feature_img)

    #--------------------------------------------#
    #   对模型的卷积核进行可视化
    #   therd_size指定保存图像的大小
    #--------------------------------------------#
    def get_conv_visual_input(self, therd_size):
        # 遍历模型中每一层网络的参数
        for name, parameters in self.model.named_parameters():
            # 如果参数的维度为4,说明是卷积层
            if parameters.ndim == 4:

                # 卷积核的个数
                kernel_number = parameters.shape[0]
                # 卷积核的通道数
                kernel_channel = parameters.shape[1]
                print(parameters.shape)
                if kernel_channel == 3:
                    dst_path = os.path.join(visual_out_path, "conv_visual", "layer_i{}".format(name))
                    makedir(dst_path)
                    # 遍历卷积核的个数
                    for n in range(kernel_number):
                        kernel = parameters[n, :, :, :]

                        feature_img = np.asarray(kernel.cpu().detach().numpy() * 255, dtype=np.uint8)
                        # 对图像进行伪彩色增强,热力图
                        # feature_img = cv2.applyColorMap(feature_img, cv2.COLORMAP_JET)
                        feature_img = np.transpose(feature_img, [1, 2, 0])
                        # 如果图片尺寸小于therd_size, 将图像resize到therd_size尺寸
                        if feature_img.shape[0] < therd_size:
                            tmp_file = os.path.join(dst_path, "channel_{}_{}.jpg".format(n, therd_size))
                            tmp_img = feature_img.copy()
                            tmp_img = cv2.resize(tmp_img, (therd_size, therd_size), interpolation=cv2.INTER_CUBIC)
                            cv2.imwrite(tmp_file, tmp_img)
                            print(tmp_img.shape)
                        else:
                            dst_file = os.path.join(dst_path, "channel_{}.jpg".format(n))
                            cv2.imwrite(dst_file, feature_img)

    #--------------------------------------------#
    #   对模型的卷积核进行可视化
    #   therd_size指定保存图像的大小
    #--------------------------------------------#
    def get_large_conv_visual(self, therd_size):
        # 遍历模型中每一层网络的参数
        for name, parameters in self.model.named_parameters():
            # 如果参数的维度为4,说明是卷积层
            if parameters.ndim == 4:

                # 卷积核的个数
                kernel_number = parameters.shape[0]
                # 卷积核的通道数
                kernel_channel = parameters.shape[1]
                # 卷积核的尺寸
                kernel_size = parameters.shape[2]
                if kernel_size == 5 or kernel_size == 7:
                    dst_path = os.path.join(visual_out_path, "conv_visual", "layer_i{}".format(name))
                    makedir(dst_path)
                    print(parameters.shape)
                    # 遍历卷积核的个数
                    for n in range(kernel_number):
                        kernel = parameters[n, :, :, :]
                        # print(kernel.shape)
                        feature_img = np.asarray(kernel.cpu().detach().numpy() * 255, dtype=np.uint8)
                        # 对图像进行伪彩色增强,热力图
                        # feature_img = cv2.applyColorMap(feature_img, cv2.COLORMAP_JET)
                        feature_img = np.transpose(feature_img, [1, 2, 0])
                        # 如果图片尺寸小于therd_size, 将图像resize到therd_size尺寸
                        if feature_img.shape[0] < therd_size:
                            tmp_file = os.path.join(dst_path, "channel_{}_{}.jpg".format(n, therd_size))
                            tmp_img = feature_img.copy()
                            tmp_img = cv2.resize(tmp_img, (therd_size, therd_size), interpolation=cv2.INTER_CUBIC)
                            cv2.imwrite(tmp_file, tmp_img)
                        else:
                            dst_file = os.path.join(dst_path, "channel_{}.jpg".format(n))
                            cv2.imwrite(dst_file, feature_img)

    #--------------------------------------------#
    #   绘制类别热力图(Class Activation, MapCAM)
    #--------------------------------------------#
    def draw_CAM(self, img_path):

        '''
        绘制 Class Activation Map
        :param img_path: str, 测试图片路径
        :param save_path: str, CAM结果保存路径
        :return:
        '''

        # 设置类别热力图的保存文件夹
        dst_path = os.path.join(visual_out_path, "cam_visual")
        makedir(dst_path)

        # 获取图片的名字
        pathList = img_path.split("/")
        imageName = pathList[-1]

        image = Image.open(img_path)
        #---------------------------------------------#
        #   对测试图片预处理
        #---------------------------------------------#
        image_data = predict_input(image, self.input_shape)
        #---------------------------------------------------------#
        #   添加batch_size维度
        #---------------------------------------------------------#
        image_data = np.expand_dims(image_data, 0)

        with torch.no_grad():
            x = torch.from_numpy(image_data)
            if self.cuda:
                x = x.cuda()

        #---------------------------------------------------------#
        #   获取模型输出的features/score
        #---------------------------------------------------------#
        bypool = False
        # for name, module in self.model._modules.items():
        for name, module in self.model.module._modules.items():
            if "avgpool" in name:
                x = module(x)
                bypool = True
                x = torch.flatten(x, 1)
            else:
                x = module(x)
                if bypool:
                    output = x
                else:
                    features = x

        #---------------------------------#
        #   为了能读取到中间梯度定义的辅助函数
        #---------------------------------#
        def extract(g):
            global features_grad
            features_grad = g

        # 预测得分最高的那一类对应的输出score
        pred = torch.argmax(output).item()
        pred_class = output[:, pred]

        #---------------------------------#
        #   读取中间features的梯度
        #---------------------------------#
        features.register_hook(extract)

        # 计算梯度
        pred_class.backward()
        # 获取梯度
        grads = features_grad

        # 计算梯度的最大池化
        pooled_grads = torch.nn.functional.adaptive_avg_pool2d(grads, (1, 1))

        # 此处batch_size默认为1, 所以去掉了第0维(batch_size维度)
        pooled_grads = pooled_grads[0]
        features = features[0]

        # 得到features的通道数
        features_channel = features.shape[0]
        for i in range(features_channel):
            features[i, ...] *= pooled_grads[i, ...]

        # 对features进行处理
        heatmap = features.cpu().detach().numpy()
        # 对features做平均池化
        heatmap = np.mean(heatmap, axis=0)
        # 进行ReLU操作
        heatmap = np.maximum(heatmap, 0)
        # 进行归一化操作
        heatmap = heatmap - np.min(heatmap)
        heatmap /= np.max(heatmap)

        # 用cv2加载原始图像
        img = cv2.imread(img_path)
        # 将热力图的大小调整为与原始图像相同
        heatmap = cv2.resize(heatmap, (img.shape[1], img.shape[0]))
        # 将热力图转换为RGB格式
        heatmap = np.uint8(255 * heatmap)
        # 将热力图应用于原始图像
        heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
        # 这里的0.3是热力图强度因子
        superimposed_img = heatmap * 0.4 + img * 0.6
        # 将图像保存到硬盘
        save_path = os.path.join(visual_out_path, "cam_visual", imageName)
        cv2.imwrite(save_path, superimposed_img)


if __name__ == "__main__":

    netVisualize = NetVisualize()
    # netVisualize.get_model_info()
    netVisualize.get_feature_output('./datasets/strongmainfire/Image546.jpg', 256)
    # netVisualize.get_conv_feature_output('./datasets/alwayslight/D06Image1980.jpg', 256)
    # netVisualize.get_conv_visual_input(256)
    # netVisualize.get_large_conv_visual(256)
    # netVisualize.get_model_metrics()
    # image = Image.open("datasets/test/cat/cat.10003.jpg")
    # netVisualize.get_feature_output(image, 256)
    # netVisualize.get_conv_visual(256)
    # imagesList = ['./datasets/none/D06Image7.jpg', './datasets/none/D06Image17895.jpg',
    #               './datasets/alwayslight/D06Image1980.jpg', './datasets/alwayslight/D06Image12295.jpg',
    #               './datasets/alwayslight/D06Image16282.jpg', './datasets/alwayslight/Image3.jpg',
    #               './datasets/weakmainfire/D06Image2606.jpg', './datasets/weakmainfire/D06Image3487.jpg',
    #               './datasets/strongmainfire/D06Image5272.jpg', './datasets/strongmainfire/Image546.jpg']
    # for i, image in enumerate(imagesList):
    #     netVisualize.draw_CAM(image)
    #     print(i + 1)


